#!/usr/bin/env python3

import rospy
from sensor_msgs.msg import Imu
import json
import rospkg
import numpy as np
import math
from tf.transformations import quaternion_from_euler, quaternion_multiply, quaternion_matrix

# 全局变量存储外参
roll = 0.0
pitch = 0.0
yaw = 0.0
x = 0.0
y = 0.0
z = 0.0

imu2_pub = None

def read_extrinsic_parameters():
    global roll, pitch, yaw, x, y, z
    try:
        rospack = rospkg.RosPack()
        package_path = rospack.get_path("livox_ros_driver2")
    except rospkg.ResourceNotFound:
        rospy.logerr("Package 'livox_ros_driver2' not found")
        return False
    
    config_file = package_path + "/config/MID360_config.json"
    
    try:
        with open(config_file, 'r') as f:
            data = json.load(f)
    except IOError:
        rospy.logerr("Cannot open config file: %s", config_file)
        return False
    except ValueError:
        rospy.logerr("Failed to parse JSON file")
        return False
    
    if "lidar_configs" not in data or not isinstance(data["lidar_configs"], list):
        rospy.logerr("Invalid lidar_configs in JSON")
        return False
    
    lidar_configs = data["lidar_configs"]
    if len(lidar_configs) == 0:
        rospy.logerr("No lidar configs available")
        return False
    
    first_config = lidar_configs[0]
    if "extrinsic_parameter" not in first_config:
        rospy.logerr("Extrinsic parameters missing")
        return False
    
    extrinsic = first_config["extrinsic_parameter"]
    try:
        roll = extrinsic["roll"]
        pitch = extrinsic["pitch"]
        yaw = extrinsic["yaw"]
        x = extrinsic["x"]
        y = extrinsic["y"]
        z = extrinsic["z"]
    except KeyError as e:
        rospy.logerr("Missing key in extrinsic: %s", str(e))
        return False
    
    return True

def imu_callback(imu_msg):
    global roll, pitch, yaw, imu2_pub
    
    # 转换欧拉角到弧度
    roll_rad = math.radians(roll)
    pitch_rad = math.radians(pitch)
    yaw_rad = math.radians(yaw)
    
    # 生成旋转四元数 (XYZ顺序)
    rotation_quat = quaternion_from_euler(roll_rad, pitch_rad, yaw_rad, axes='sxyz')
    
    # 获取3x3旋转矩阵
    R = quaternion_matrix(rotation_quat)[:3, :3]
    
    # 处理线性加速度
    acc = np.array([imu_msg.linear_acceleration.x,
                    imu_msg.linear_acceleration.y,
                    imu_msg.linear_acceleration.z])
    rotated_acc = R.dot(acc)
    
    # 处理角速度
    angular = np.array([imu_msg.angular_velocity.x,
                        imu_msg.angular_velocity.y,
                        imu_msg.angular_velocity.z])
    rotated_angular = R.dot(angular)
    
    # 处理方向四元数
    q_orig = np.array([imu_msg.orientation.x,
                       imu_msg.orientation.y,
                       imu_msg.orientation.z,
                       imu_msg.orientation.w])
    
    q_rotated = quaternion_multiply(rotation_quat, q_orig)
    
    # 构造新的IMU消息
    imu2_msg = Imu()
    imu2_msg.header = imu_msg.header
    
    # 填充加速度
    imu2_msg.linear_acceleration.x = rotated_acc[0]
    imu2_msg.linear_acceleration.y = rotated_acc[1]
    imu2_msg.linear_acceleration.z = rotated_acc[2]
    
    # 填充角速度
    imu2_msg.angular_velocity.x = rotated_angular[0]
    imu2_msg.angular_velocity.y = rotated_angular[1]
    imu2_msg.angular_velocity.z = rotated_angular[2]
    
    # 填充方向四元数
    imu2_msg.orientation.x = q_rotated[0]
    imu2_msg.orientation.y = q_rotated[1]
    imu2_msg.orientation.z = q_rotated[2]
    imu2_msg.orientation.w = q_rotated[3]
    
    # 发布消息
    imu2_pub.publish(imu2_msg)

def main():
    global imu2_pub
    rospy.init_node('imu_processor')
    
    if not read_extrinsic_parameters():
        rospy.logerr("Failed to initialize extrinsic parameters")
        return
    
    
    imu2_pub = rospy.Publisher('/livox/imu2', Imu, queue_size=5)
    rospy.Subscriber('/livox/imu', Imu, imu_callback, queue_size=5)
    
    rospy.spin()

if __name__ == '__main__':
    main()